import traffic_pb2, sys, os
from datetime import datetime

# Reads the traffic matrices
class TrafficSnapshotReader(object):
	def __init__(self, dcn_name, sharded_file_directory="/Users/minyee/Desktop/topology_engineering_data/traffic_modes/sharded_traffic_files/"):
		self.dcn_name = dcn_name
		self.file_directory = sharded_file_directory
		return
	
	def _get_blocks(self, traffic_snapshots):
		blocks = set()
		for traffic_snapshot in traffic_snapshots:
			for record in traffic_snapshot.records:
				blocks.add(record.source)
				blocks.add(record.destination)
		return blocks

	def _load_traffic_history(self):
		idx = 0
		sharded_fname = self.file_directory + self.dcn_name + "_{}".format(idx)
		traffic_history = traffic_pb2.TrafficHistory()
		all_blocks_set = None
		traffic_snapshots = []
		#while os.path.isfile(sharded_fname) and idx <=5:
		
		#while os.path.isfile(sharded_fname) and idx <= 3:
		
		#while os.path.isfile(sharded_fname) and idx <=5:	
		#while os.path.isfile(sharded_fname) and idx <= 1:
		while os.path.isfile(sharded_fname) and idx <=9:	
			with open(sharded_fname, 'r') as file:
				traffic_history.ParseFromString(file.read())
			traffic_snapshots += traffic_history.snapshots
			idx += 1
			sharded_fname = self.file_directory + self.dcn_name + "_{}".format(idx)
		all_blocks_set = self._get_blocks(traffic_history.snapshots)
		return all_blocks_set , traffic_snapshots


	## External method
	# reads all the traffic matrices, and returns a list of traffic snapshots and a list of block names
	def get_traffic_snapshots(self):
		all_blocks_set , traffic_snapshots = self._load_traffic_history()
		return list(all_blocks_set), traffic_snapshots
